package com.lunatic.assistant.controller;

import com.lunatic.assistant.model.ActorFilms;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;
import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@RestController
public class AssistantController {

    @Autowired
    private ChatClient.Builder chatClient;

    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private EmbeddingModel embeddingModel;


    @Autowired
    private ChatModel chatModel;


    /**
     * 同步响应
     *
     * @param s
     * @return
     */
    @RequestMapping("demo1")
    public String demo1(@RequestParam("s") String s) {
        return chatClient.build().prompt(s).call().content();
    }

    /**
     * 流式响应
     *
     * @param s
     * @return
     */
    @RequestMapping("demo2")
    public Flux<String> demo2(@RequestParam("s") String s) {
        return chatClient.build().prompt(s).stream().content();
    }

    /**
     * chatClient demo
     * 默认消息
     *
     * @param s
     * @return
     */
    @RequestMapping("demo3")
    public Flux<String> demo3(@RequestParam("s") String s) {
        return chatClient.defaultSystem("You are a friendly chat bot that answers question in the voice of a Pirate")
                .build()
                .prompt(s)
                .stream()
                .content();
    }

    /**
     * 使用模版
     *
     * @param s
     * @param voice
     * @return
     */
    @RequestMapping("demo33")
    public Flux<String> demo33(@RequestParam("s") String s, @RequestParam("voice") String voice) {
        return chatClient.defaultSystem("You are a friendly chat bot that answers question in the voice of a {voice}")
                .build()
                .prompt(s)
                .system(ps -> ps.param("voice", voice))
                .stream()
                .content();
    }


    /**
     * chatClient demo
     * 流式响应，设置记忆对话的大小
     *
     * @param s
     * @returna
     */
    @RequestMapping("demo4")
    public Flux<String> demo4(@RequestParam("s") String s) {
        return chatClient
                .defaultAdvisors(
                        new PromptChatMemoryAdvisor(chatMemory),
                        new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().build())
                )
                .build()
                .prompt()
                .user(s)
                .advisors(a -> a
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, 1)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .stream().content();
    }

    /**
     * chatClient demo
     * 流式响应，预设文本
     *
     * @param s
     * @returna
     */
    @RequestMapping("demo5")
    public Flux<String> demo5(@RequestParam("s") String s) {
        return chatClient.defaultSystem("今天是{datetime}")
                .build()
                .prompt()
                .system(p -> p.param("datetime", LocalDate.now()))
                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().build()))
                .user(s)
                .stream()
                .content();
    }


    /**
     * chatClient demo
     * 日志记录
     *
     * @returna
     */
    @RequestMapping("demo6")
    public String demo6(@RequestParam("s") String s) {
        SimpleLoggerAdvisor customLogger = new SimpleLoggerAdvisor(
                request -> "Custom request: " + request.userText(),
                response -> "Custom response: " + response.getResult(),
                1
        );
        return chatClient
                .build()
                .prompt()
                .advisors(customLogger)
                .user(s)
                .call()
                .content();
    }

    @GetMapping("/demo7")
    public Map embed(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(s));
        return Map.of("embedding", embeddingResponse);
    }


    @GetMapping("/demo8")
    public Flux<String> demo8(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return chatClient.defaultFunctions("weatherFunction", "getTemper", "orderTik")
                .build().prompt().user(s).stream().content();
    }

    @GetMapping("/demo9")
    public Flux<String> demo9(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return chatClient.defaultSystem("你将作为一名机器人产品的专家，对于用户的使用需求作出解答")
                .build()
                .prompt()
                .user(s)
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .stream()
                .content();
    }

    @GetMapping("/demo10")
    public String demo10(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {

        ChatClient.Builder builder = chatClient.defaultSystem("你是一位专业的室内设计顾问，精通各种装修风格、材料选择和空间布局。请基于提供的参考资料，为用户提供专业、详细且实用的建议。在回答时，请注意：\\n\" +\n" +
                "                \"1. 准确理解用户的具体需求\\n\" +\n" +
                "                \"2. 结合参考资料中的实际案例\\n\" +\n" +
                "                \"3. 提供专业的设计理念和原理解释\\n\" +\n" +
                "                \"4. 考虑实用性、美观性和成本效益\\n\" +\n" +
                "                \"5. 如有需要，可以提供替代方案");

        // 构建查询扩展器, 用于生成多个相关的查询变体，以获得更全面的搜索结果
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(builder)
                .includeOriginal(false) // 不包含原始查询
                .numberOfQueries(10) // 生成3个查询变体
                .build();
        // 执行查询扩展, 将原始问题"请提供几种推荐的装修风格?"扩展成多个相关查询
        List<Query> queries = queryExpander.expand(new Query("请提供几种推荐的装修风格"));
        return queries.toString();
    }

    @GetMapping("/demo11")
    public void demo11(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        ChatClient.Builder builder = chatClient;
        // 创建一个模拟用户学习AI的查询场景
        Query query = new Query("我正在学习人工智能，什么是大语言模型？");
        // 创建查询重写转换器
        QueryTransformer queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
        // 执行查询重写
        Query transformedQuery = queryTransformer.transform(query);
        // 输出重写后的查询
        System.out.println(transformedQuery.text());
    }

    @GetMapping("/demo12")
    public void demo12(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        ChatClient.Builder builder = chatClient;
        // QueryTransformer用于将带有上下文的查询转换为完整的独立查询
        QueryTransformer queryTransformer = CompressionQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
        // 这个例子模拟了一个房地产咨询场景，用户先问小区位置，再问房价
        Query query = Query.builder()
                .text("那这个小区的二手房均价是多少?")  // 当前用户的提问
                .history(new UserMessage("深圳市南山区的碧海湾小区在哪里?"),  // 历史对话中用户的问题
                        new AssistantMessage("碧海湾小区位于深圳市南山区后海中心区，临近后海地铁站。"))  // AI的回答
                .build();        // 将模糊的代词引用（"这个小区"）转换为明确的实体名称（"碧海湾小区"）
        Query transformedQuery = queryTransformer.transform(query);
        // 输出重写后的查询
        System.out.println(transformedQuery.text());
    }


    @GetMapping("/demo13")
    public ActorFilms demo13(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return ChatClient.create(chatModel).prompt()
                .user(u -> u.text("Generate the filmography of 5 movies for {actor}.")
                        .param("actor", "Tom Hanks"))
                .call()
                .entity(ActorFilms.class);
    }

    @GetMapping("/demo14")
    public List<ActorFilms> demo14(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return ChatClient.create(chatModel).prompt()
                .user("Generate the filmography of 5 movies for Tom Hanks and Bill Murray.")
                .call()
                .entity(new ParameterizedTypeReference<>() {
                });
    }

    @GetMapping("/demo15")
    public Map<String, Object> demo15(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return ChatClient.create(chatModel).prompt()
                .user(u -> u.text("Provide me a List of {subject}")
                        .param("subject", "an array of numbers from 1 to 9 under their key name 'numbers'"))
                .call()
                .entity(new ParameterizedTypeReference<>() {
                });
    }

    @GetMapping("/demo16")
    public List<String> demo16(@RequestParam(value = "s", defaultValue = "Tell me a joke") String s) {
        return ChatClient.create(chatModel).prompt()
                .user(u -> u.text("List five {subject}")
                        .param("subject", "ice cream flavors"))
                .call()
                .entity(new ListOutputConverter(new DefaultConversionService()));
    }

}
